import json
import httpx
from typing import Dict, Any, Union, AsyncGenerator
from app.core.coze_config import coze_settings
from app.models.coze_schemas import CozeMessage, CozeApiRequest

class CozeService:
    def __init__(self):
        self.api_url = coze_settings.COZE_API_URL
        self.headers = {
            "Authorization": f"Bearer {coze_settings.COZE_API_KEY}",
            "Content-Type": "application/json"
        }

    async def send_message_stream(self, content: str) -> AsyncGenerator[str, None]:
        message = CozeMessage(content=content)
        request_data = CozeApiRequest(
            bot_id=coze_settings.COZE_BOT_ID,
            user_id=coze_settings.COZE_USER_ID,
            stream=True,
            auto_save_history=coze_settings.AUTO_SAVE_HISTORY,
            additional_messages=[message]
        )

        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                self.api_url,
                headers=self.headers,
                json=request_data.model_dump(exclude_none=True),
                timeout=30.0
            ) as response:
                if response.status_code != 200:
                    error_data = await response.json()
                    raise Exception(f"API Error: {error_data.get('msg')}")

                async for line in response.aiter_lines():
                    if line.strip():
                        try:
                            data = json.loads(line)
                            if data.get('code') != 0:
                                yield f"event: error\ndata: {json.dumps(data)}\n\n"
                                break
                            yield f"event: message\ndata: {json.dumps(data)}\n\n"
                        except json.JSONDecodeError:
                            continue
                yield "event: done\ndata: {\"status\": \"completed\"}\n\n"

    async def send_message_sync(self, content: str) -> Dict[str, Any]:
        message = CozeMessage(content=content)
        request_data = CozeApiRequest(
            bot_id=coze_settings.COZE_BOT_ID,
            user_id=coze_settings.COZE_USER_ID,
            stream=False,
            auto_save_history=coze_settings.AUTO_SAVE_HISTORY,
            additional_messages=[message]
        )

        async with httpx.AsyncClient() as client:
            response = await client.post(
                self.api_url,
                headers=self.headers,
                json=request_data.model_dump(exclude_none=True),
                timeout=30.0
            )
            response_data = response.json()
            if response.status_code != 200 or response_data.get('code') != 0:
                raise Exception(f"API Error: {response_data.get('msg')}")
            return response_data


